import os
import cv2
import numpy as np
from flask import Flask, request, render_template, jsonify
from werkzeug.utils import secure_filename
from paddleocr import PaddleOCR

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'static/uploads'
app.config['ALLOWED_EXTENSIONS'] = {'png', 'jpg', 'jpeg'}

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

# 初始化OCR模型
ocr = PaddleOCR(use_angle_cls=True, lang="ch", use_gpu=False)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/recognize', methods=['POST'])
def recognize():
    if 'file' not in request.files:
        return jsonify({'error': '没有上传文件'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '未选择文件'}), 400
    
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        # 使用PaddleOCR进行识别
        result = ocr.ocr(filepath, cls=True)
        
        # 处理识别结果
        ocr_results = []
        for line in result[0]:
            text = line[1][0]  # 提取识别的文本
            confidence = line[1][1]  # 提取置信度
            ocr_results.append({
                'text': text,
                'confidence': float(confidence)
            })
        
        # 在图像上绘制识别结果
        img = cv2.imread(filepath)
        for line in result[0]:
            # 获取文本框坐标
            points = line[0]
            points = np.array(points).astype(np.int32)
            
            # 绘制边框
            cv2.polylines(img, [points], True, (0, 255, 0), 2)
            
            # 添加文本
            text = line[1][0]
            cv2.putText(img, text, (points[0][0], points[0][1] - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1)
        
        # 保存标注后的图像
        output_filename = 'result_' + filename
        output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
        cv2.imwrite(output_path, img)
        
        return jsonify({
            'result': ocr_results,
            'original_image': '/static/uploads/' + filename,
            'annotated_image': '/static/uploads/' + output_filename
        })
    
    return jsonify({'error': '不支持的文件类型'}), 400

if __name__ == '__main__':
    app.run(debug=True) 